import argparse
from pathlib import Path
import cv2
import numpy as np
import os

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="")
    parser.add_argument('--origin-path', default=r"C:\Users\ouyangwj\Desktop\images2_person\003_0060.jpg", type=str)
    parser.add_argument('--target-path', default=r"C:\Users\ouyangwj\Desktop\images_mask", type=str)
    return parser.parse_args()

def Draw_Mask(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        param['lbutton_state'] = 1
        param['init_point'] = (x, y)
        param['previous_image_states'].append(param['image_mask'].copy())
        if flags & cv2.EVENT_FLAG_SHIFTKEY:
            param['draw_mode'] = 'rectangle'
        else:
            param['draw_mode'] = 'freeform'

    elif event == cv2.EVENT_MOUSEMOVE:
        if param['lbutton_state'] == 1:
            if param['draw_mode'] == 'rectangle':
                # Display a temporary rectangle on a copy of the image
                temp_image = param['previous_image_states'][-1].copy()
                cv2.rectangle(temp_image, param['init_point'], (x, y), (0, 0, 0), -1)
                cv2.imshow(f"{param['image_path'].name}", cv2.bitwise_and(param['image'], temp_image))
            elif param['draw_mode'] == 'freeform':
                cv2.line(param['image_mask'], param['init_point'], (x, y), (0, 0, 0), param["draw_size"])
                param['init_point'] = (x, y)
                cv2.imshow(f"{param['image_path'].name}", cv2.bitwise_and(param['image'], param['image_mask']))

    elif event == cv2.EVENT_LBUTTONUP:
        param['lbutton_state'] = 0
        if param['draw_mode'] == 'rectangle':
            # Fix the rectangle on the actual mask image
            cv2.rectangle(param['image_mask'], param['init_point'], (x, y), (0, 0, 0), -1)
        else:
            cv2.line(param['image_mask'], param['init_point'], (x, y), (0, 0, 0), param["draw_size"])
        cv2.imwrite(f"{param['target_img_dir']}{os.sep}{param['image_path'].stem}_mask.jpg", param['image_mask'])
        print(f"Mask saved at {param['target_img_dir']}{os.sep}{param['image_path'].stem}_mask.jpg")

    elif event == cv2.EVENT_RBUTTONDOWN:
        if param['previous_image_states']:
            param['image_mask'] = param['previous_image_states'].pop()
            cv2.imshow(f"{param['image_path'].name}", cv2.bitwise_and(param['image'], param['image_mask']))

def Get_Draw_Size(x, params):
    params['draw_size'] = x

def main():
    args = Parse_Arguments()
    origin_path = Path(args.origin_path.replace("\\", "/"))
    target_path = Path(args.target_path.replace("\\", "/"))
    image = cv2.imread(str(origin_path))
    if image is None:
        print(f"Error: Unable to load image at {origin_path}. Please check the file path.")
        return
    origin_path_mask = Path(str(origin_path.parent) + '/' + origin_path.stem + "_mask.jpg")
    if origin_path_mask.exists():
        mask = cv2.imread(str(origin_path_mask))
    else:
        mask = np.full((image.shape[0], image.shape[1], 3), 255, np.uint8)

    params = {
        'draw_size': 20,
        'init_point': (0, 0),
        'lbutton_state': 0,
        'draw_mode': 'freeform',
        'previous_image_states': [mask.copy()],
        'image_mask': mask,
        'image': image,
        'image_path': origin_path,
        'target_img_dir': target_path
    }
    cv2.namedWindow(f'{origin_path.name}', cv2.WINDOW_NORMAL)
    cv2.moveWindow(f'{origin_path.name}', 600, 250)
    cv2.resizeWindow(f'{origin_path.name}', 1500, 980)
    cv2.createTrackbar('draw_size', f'{origin_path.name}', 10, 255, lambda x: Get_Draw_Size(x, params))
    cv2.setMouseCallback(f'{origin_path.name}', Draw_Mask, params)
    cv2.imshow(f'{origin_path.name}', image)
    cv2.waitKeyEx(0)

if __name__ == "__main__":
    main()
